/*
 * Copyright 2008 Carnegie Mellon University
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *  
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, 
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package edu.cmu.lti.sepia.data.db;

import java.io.IOException;
import java.io.OutputStream;
import java.io.Writer;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import jxl.Workbook;
import jxl.WorkbookSettings;
import jxl.write.Label;
import jxl.write.WritableSheet;
import jxl.write.WritableWorkbook;
import jxl.write.WriteException;
import jxl.write.biff.RowsExceededException;
import edu.cmu.lti.sepia.infra.eval.model.Pool;
import edu.cmu.lti.sepia.infra.eval.model.QAJudgmentResult;
import edu.cmu.lti.sepia.infra.eval.model.QAJudgmentResults;
import edu.cmu.lti.sepia.infra.eval.model.QAJudgmentUnit;
import edu.cmu.lti.sepia.infra.topic.model.Nugget;
import edu.cmu.lti.sepia.infra.topic.model.Topic;

/**
 * 
 * Manages QA Judgment Result object in a DB.
 * 
 * @see QAJudgmentResult
 * @author Hideki Shima
 *
 */
public class QAJudgmentResultDAO {
  
  public static QAJudgmentResults findResult( SQL sql, String table, QAJudgmentUnit unit ) 
  throws SQLException {
    int judgmentId = unit.getJudgmentId();
//    if (judgmentId<=0) {
//      unit = QAJudgmentUnitDAO.findJudgmentUnitByTopicAndDocument(sql, table, unit.getTopicId(), unit.getDocumentId());
//      judgmentId = unit.getJudgmentId();
//    } else if (unit.getSystemResponse()==null || unit.getTopicId()==null ||
//            unit.getSystemResponse().length()==0 || unit.getTopicId().length()==0 ) {
//      //Problem!!
//      unit = QAJudgmentUnitDAO.findJudgmentUnitByTopicAndDocument(sql, table, unit.getTopicId(), unit.getDocumentId());
//    }
    String query = Query.generate( Query.FIND_QA_JUDGMENT_RESULT, table );
    PreparedStatement ps = sql.getConnection().prepareStatement(query);
    ResultSet rs = null;
    QAJudgmentResults result = new QAJudgmentResults();
    
    ps.setInt(1, judgmentId);
    rs = sql.executeSelect(ps);
    if ( rs == null ) System.out.println( ps );
    while ( rs.next() ) {
      QAJudgmentResult r = rsToObject(rs);
      r.setUnit( unit );
      result.add( r );
    }
    rs.close();
    ps.close();
    return result;
  }
  
  public static QAJudgmentResults findResultsByTopicAndPool( SQL sql, String table, String topicId, int poolId )
  throws SQLException {
    QAJudgmentResults results = new QAJudgmentResults();
    List<QAJudgmentUnit> units = QAJudgmentUnitDAO.findJudgmentUnitsByTopicAndPool(sql, table, topicId, poolId);
    for (QAJudgmentUnit u : units) {
      QAJudgmentResults resultsByAssessors = findResult(sql, table, u);
      if ( resultsByAssessors.size()>0 ) { 
        results.addAll(resultsByAssessors);
      } else {// not judged yet
        results.add(new QAJudgmentResult(u, -1, -1, "", ""));
      }
    }
    return results;
  }
  
  public static QAJudgmentResults findResultsByTopic( SQL sql, String table, String topicId )
  throws SQLException {
    QAJudgmentResults results = new QAJudgmentResults();
    List<Pool> pools = PoolDAO.getPools(sql, table, false);
    for (Pool p : pools) {
      QAJudgmentResults resultsByPool = findResultsByTopicAndPool(sql, table, topicId, p.getPoolId());
      results.addAll(resultsByPool);
    }
    return results;
  }
  
  public static void updateResult( SQL sql, String table, 
    int judgmentId, int nuggetId, int userId, int matched ) throws SQLException {
    String query = Query.generate( Query.UPDATE_QA_JUDGMENT_RESULT, table );
    PreparedStatement ps = sql.getConnection().prepareStatement(query);
    ps.setInt(1, judgmentId);
    ps.setInt(2, nuggetId);
    ps.setInt(3, userId);
    ps.setInt(4, matched);
    ps.setInt(5, matched);
    ps.executeUpdate();
    ps.close();
  }
  
  public static void exportJudgments( SQL sql, String table, String trgLang, Writer out ) 
  throws SQLException, IOException {
    List<Topic> topics = TopicDAO.getGoldStandard( sql, table, -1, null );
    TopicDAO.removeSimilarNuggets(topics);
    TopicDAO.generateExternalIds(topics,trgLang,4);
    TopicDAO.removeUnselected(topics);
    
    for ( Topic t : topics ) {
      StringBuilder sb = new StringBuilder();
      String tid = t.getExternalId();
      
      List<QAJudgmentUnit> units = QAJudgmentUnitDAO.findJudgmentUnitsByTopic(sql, table, tid);      
      List<QAJudgmentResult> results = QAJudgmentResultDAO.findResultsByTopic(sql, table, tid);
      
      sb.append("Submission for topic "+tid+" with items numbered:\n");
      for ( int i=0; i<units.size(); i++ ) {
        QAJudgmentUnit u = units.get(i);
        sb.append(tid+"\t");
        sb.append((i+1)+"\t");
        sb.append(u.getSystemResponse()+"\n");
      }
      //QAJudgmentUnits units = QAJudgmentUnitDAO.findJudgmentUnitsByTopicAndPool(sql, table, topicId, poolId)
      sb.append("\n");
      sb.append("Nuggets assigned to items for topic "+tid+":\n");
      for ( int i=0; i<units.size(); i++ ) {
        QAJudgmentUnit u = units.get(i);
        SortedSet<Integer> set = new TreeSet<Integer>();
        for ( QAJudgmentResult r : results ) { 
          if ( u.getJudgmentId() == r.getJudgmentId() ) {
            if (r.getJudgment().equals("1") && t.getNuggets()!=null) {
              for ( int j=0; j<t.getNuggets().size(); j++ ) {
                Nugget n = t.getNuggets().get(j);
                if (n.getNuggetId() == r.getNuggetId()) {
                  set.add( (j+1) );
                }
              }
            }
          }
        }
        if (set.size()>0) {
          sb.append( tid+"\t"+(i+1)+"\t" );
          for ( Integer j : set ) {
            sb.append( j+" " );
          }
          sb.append("\n");
        }
      }
      sb.append("\n\n\n");
      out.write(sb.toString());
    }
  }
  
  // Since the qrels are huge, don't store them into objects
  public static void exportJudgmentsInExcel( SQL sql, String table, String lang, OutputStream os ) 
  throws SQLException, IOException, RowsExceededException, WriteException {
    List<Topic> topics = TopicDAO.getGoldStandard( sql, table, -1, null );
    TopicDAO.removeSimilarNuggets(topics);
    TopicDAO.generateExternalIds(topics,lang,4);
    TopicDAO.removeUnselected(topics);
    WorkbookSettings ws = new WorkbookSettings();
    if (lang.equals("JA")) {
      ws.setLocale(Locale.JAPANESE);
      //ws.setEncoding("Windows-31J");
    } else if (lang.equals("CS")) {
      ws.setLocale(Locale.SIMPLIFIED_CHINESE);
    } else if (lang.equals("CT")) {
      ws.setLocale(Locale.TRADITIONAL_CHINESE);
    }

    WritableWorkbook workbook = Workbook.createWorkbook(os, ws);
    
    for ( int ti=0; ti<topics.size(); ti++ ) {
      Topic t = topics.get(ti);
      if (t.getNuggets()==null || t.getNuggets().size()==0) continue;
      String tid = t.getExternalId();
      
      List<QAJudgmentUnit> units = QAJudgmentUnitDAO.findJudgmentUnitsByTopic(sql, table, tid);      
      List<QAJudgmentResult> results = QAJudgmentResultDAO.findResultsByTopic(sql, table, tid);
      
      WritableSheet sheet = workbook.createSheet(t.getExternalId(), ti);
      
      // write system responses
      for ( int yi=0; yi<units.size(); yi++ ) {
        QAJudgmentUnit u = units.get(yi);
        sheet.addCell( new Label(0,yi+2,"System Response "+(yi+1)) );
        sheet.addCell( new Label(1,yi+2,u.getSystemResponse()) );
      }
      
      // write nuggets
      for ( int xi=0; xi<t.getNuggets().size(); xi++ ) {
        sheet.addCell( new Label(xi+2,0,"Nugget "+(xi+1)) );
        sheet.addCell( new Label(xi+2,1,t.getNuggets().get(xi).getText()) );
      }
      
      // init judgments
      for (int xi=0; xi<t.getNuggets().size(); xi++) {
        for ( int yi=0; yi<units.size(); yi++ ) {
          sheet.addCell( new Label(xi+2,yi+2,"-") );
        }
      }
      
      // write judgments
      for ( int i=0; i<units.size(); i++ ) {
        QAJudgmentUnit u = units.get(i);
        for ( QAJudgmentResult r : results ) { 
          if ( u.getJudgmentId() == r.getJudgmentId() ) {
            if (r.getJudgment().equals("1") && t.getNuggets()!=null) {
              for ( int j=0; j<t.getNuggets().size(); j++ ) {
                Nugget n = t.getNuggets().get(j);
                if (n.getNuggetId() == r.getNuggetId()) {
                  sheet.addCell( new Label(j+2,i+2,"x") );
                }
              }
            }
          }
        }
      }
    }
    workbook.write();
    workbook.close();
  }
  
  private static QAJudgmentResult rsToObject( ResultSet rs ) {
    try {
      QAJudgmentResult r = new QAJudgmentResult( rs.getInt(1), rs.getInt(2),
              rs.getInt(3), rs.getString(4));
      return r;
    } catch ( Exception e ) {
      e.printStackTrace();
      return null;
    }
  }
  
  public static JSONObject getQAJudgmentResults( List<QAJudgmentResult> results ) throws JSONException {

    Map<String,JSONArray> nuggetMap = new LinkedHashMap<String,JSONArray>();
    
    for (QAJudgmentResult result : results) {
      String key = result.getJudgmentId()+":"+result.getSystemResponse();
      JSONArray ja = nuggetMap.get(key);
      if (ja==null) {
        ja = new JSONArray();
      }
      JSONObject jNugget = new JSONObject();
      jNugget.put("nuggetId", result.getNuggetId());
      jNugget.put("matched", result.getJudgment());
      ja.put(jNugget);
      nuggetMap.put(key, ja);
    }
    
    JSONObject json = new JSONObject();
    JSONArray jsonArray = new JSONArray();
    
    for (String key : nuggetMap.keySet()) {
      JSONObject j = new JSONObject();
      String[] items = key.split(":");
      j.put("judgmentId", items[0]);
      j.put("systemResponse", items[1]);
      j.put("nuggets", nuggetMap.get(key));
      jsonArray.put(j);
    }
    json.put("qa_judgments", jsonArray);
    return json;
  }
}
